import numpy as np
from scipy.ndimage import label, find_objects
from easy_io import read_pkl_file, write_pkl_file
from utils import groupby
from itertools import chain


def _merge_a_scan(candidates):
    bboxes = np.asarray([c['bbox'] for c in candidates])
    shape = np.max(bboxes[:, -3:], axis=0)
    labels = np.zeros(shape, dtype='bool')
    slices = (tuple(slice(i, j) for i, j in zip(bbox[:3], bbox[-3:])) for bbox in bboxes)
    for s in slices:
        labels[s] = True
    labels, _ = label(labels)
    object_slices = find_objects(labels)
    bboxes = (tuple(s.start for s in t) + tuple(s.stop for s in t) for t in object_slices)
    candidates = tuple(dict(bbox=bbox) for bbox in bboxes)
    return candidates


def merge(candidates):
    candidate_dict = groupby(candidates, lambda c: c['scanid'])
    candidates = []
    for scanid, clst in candidate_dict.items():
        source = clst[0]['source']
        clst = _merge_a_scan(clst)
        for c in clst:
            c.update(scanid=scanid, source=source)
        candidates.append(clst)
    candidates = tuple(chain.from_iterable(candidates))
    return candidates


def main(infile, outfile):
    candidates = read_pkl_file(infile)
    candidates = merge(candidates)
    write_pkl_file(outfile, candidates)


if __name__ == '__main__':
    main(
        infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v2.pkl',
        outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v3.pkl',
    )
